// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/expression/vector_function.h>

namespace kumo::pollux::functions {
    namespace {
        class MapKeyValueFunction : public exec::VectorFunction {
        public:
            void apply(
                const SelectivityVector &rows,
                std::vector<VectorPtr> &args,
                const TypePtr &outputType,
                exec::EvalCtx &context,
                VectorPtr &result) const override {
                auto &arg = args[0];

                VectorPtr localResult;

                // Input can be constant or flat.
                if (arg->is_constant_encoding()) {
                    auto *constantMap = arg->as<ConstantVector<ComplexType> >();
                    const auto &flatMap = constantMap->value_vector();
                    const auto flatIndex = constantMap->index();

                    exec::LocalSingleRow singleRow(context, flatIndex);
                    localResult = applyFlat(*singleRow, flatMap, context);
                    localResult =
                            BaseVector::wrap_in_constant(rows.end(), flatIndex, localResult);
                } else {
                    localResult = applyFlat(rows, arg, context);
                }

                context.moveOrCopyResult(localResult, rows, result);
            }

        protected:
            explicit MapKeyValueFunction(const std::string &name) : name_(name) {
            }

            virtual VectorPtr applyFlat(
                const SelectivityVector &rows,
                const VectorPtr &arg,
                exec::EvalCtx &context) const = 0;

        private:
            const std::string name_;
        };

        class MapKeysFunction : public MapKeyValueFunction {
        public:
            MapKeysFunction() : MapKeyValueFunction("map_keys") {
            }

            VectorPtr applyFlat(
                const SelectivityVector &rows,
                const VectorPtr &arg,
                exec::EvalCtx &context) const override {
                POLLUX_CHECK(
                    arg->type_kind() == TypeKind::MAP,
                    "Unsupported type for map_keys function {}",
                    mapTypeKindToName(arg->type_kind()));

                auto map_vector = arg->as<MapVector>();
                auto mapKeys = map_vector->mapKeys();
                return std::make_shared<ArrayVector>(
                    context.pool(),
                    ARRAY(mapKeys->type()),
                    map_vector->nulls(),
                    rows.end(),
                    map_vector->offsets(),
                    map_vector->sizes(),
                    mapKeys,
                    map_vector->get_null_count());
            }

            static std::vector<std::shared_ptr<exec::FunctionSignature> > signatures() {
                // map(K,V) -> array(K)
                return {
                    exec::FunctionSignatureBuilder()
                    .typeVariable("K")
                    .typeVariable("V")
                    .returnType("array(K)")
                    .argumentType("map(K,V)")
                    .build()
                };
            }
        };

        class MapValuesFunction : public MapKeyValueFunction {
        public:
            MapValuesFunction() : MapKeyValueFunction("map_values") {
            }

            VectorPtr applyFlat(
                const SelectivityVector &rows,
                const VectorPtr &arg,
                exec::EvalCtx &context) const override {
                POLLUX_CHECK(
                    arg->type_kind() == TypeKind::MAP,
                    "Unsupported type for map_values function {}",
                    mapTypeKindToName(arg->type_kind()));

                auto map_vector = arg->as<MapVector>();
                auto mapValues = map_vector->mapValues();
                return std::make_shared<ArrayVector>(
                    context.pool(),
                    ARRAY(mapValues->type()),
                    map_vector->nulls(),
                    rows.end(),
                    map_vector->offsets(),
                    map_vector->sizes(),
                    mapValues,
                    map_vector->get_null_count());
            }

            static std::vector<std::shared_ptr<exec::FunctionSignature> > signatures() {
                // map(K,V) -> array(V)
                return {
                    exec::FunctionSignatureBuilder()
                    .typeVariable("K")
                    .typeVariable("V")
                    .returnType("array(V)")
                    .argumentType("map(K,V)")
                    .build()
                };
            }
        };
    } // namespace

    POLLUX_DECLARE_VECTOR_FUNCTION(
        udf_map_keys,
        MapKeysFunction::signatures(),
        std::make_unique<MapKeysFunction>());

    POLLUX_DECLARE_VECTOR_FUNCTION(
        udf_map_values,
        MapValuesFunction::signatures(),
        std::make_unique<MapValuesFunction>());
} // namespace kumo::pollux::functions
